"""
文本问答智能体
专门处理文本问答任务的智能体
"""

from typing import Dict, Any, List, Optional
import asyncio

from apps.agents.base_agent import BaseAgent, AgentType, AgentStatus
from apps.text_qa.tools import text_qa_tools
from utils.logger import logger


class TextQAAgent(BaseAgent):
    """文本问答智能体"""
    
    def __init__(self, agent_id: str = "text_qa_agent"):
        super().__init__(
            agent_id=agent_id,
            agent_type=AgentType.TEXT_QA,
            name="文本问答智能体"
        )
        
        # 默认参数
        self.default_similarity_threshold = 0.6
        self.default_max_results = 5
        
        logger.info(f"文本问答智能体初始化完成: {self.agent_id}")
    
    async def validate_input(self, input_data: Dict[str, Any]) -> bool:
        """验证输入数据"""
        required_fields = ["query"]
        
        for field in required_fields:
            if field not in input_data:
                logger.error(f"缺少必填字段: {field}")
                return False
        
        query = input_data["query"]
        if not query or not isinstance(query, str) or len(query.strip()) == 0:
            logger.error("查询内容无效")
            return False
        
        return True
    
    async def process_task(self, task_data: Dict[str, Any]) -> Dict[str, Any]:
        """处理文本问答任务"""
        query = task_data["query"]
        
        logger.info(f"开始处理文本问答任务: {query}")
        
        try:
            # 1. 关键词提取
            keywords = text_qa_tools.extract_keywords(query)
            logger.info(f"提取关键词: {keywords}")
            
            # 2. 知识库检索（这里模拟，实际应该从知识库获取）
            retrieved_texts = await self._retrieve_from_knowledge_base(
                query, keywords, task_data
            )
            
            if not retrieved_texts:
                logger.warning("未检索到相关内容")
                return {
                    "query": query,
                    "answer": "抱歉，没有找到相关信息来回答您的问题。",
                    "sources": [],
                    "confidence": 0.0
                }
            
            # 3. 生成答案
            use_model = task_data.get("use_model", True)
            similarity_threshold = task_data.get("similarity_threshold", self.default_similarity_threshold)
            
            answer = await text_qa_tools.generate_answer_async(
                query=query,
                retrieved_texts=retrieved_texts,
                similarity_threshold=similarity_threshold,
                use_model=use_model
            )
            
            if not answer:
                logger.warning("答案生成失败")
                return {
                    "query": query,
                    "answer": "抱歉，无法生成答案。",
                    "sources": retrieved_texts,
                    "confidence": 0.0
                }
            
            # 4. 计算置信度
            confidence = self._calculate_confidence(retrieved_texts, similarity_threshold)
            
            result = {
                "query": query,
                "answer": answer,
                "sources": retrieved_texts,
                "confidence": confidence,
                "keywords": keywords,
                "model_used": use_model and text_qa_tools.model_client is not None
            }
            
            logger.info(f"文本问答任务完成: {query}, 置信度: {confidence}")
            return result
            
        except Exception as e:
            logger.error(f"文本问答任务失败: {str(e)}")
            raise
    
    async def _retrieve_from_knowledge_base(self, query: str, keywords: List[str], 
                                          task_data: Dict[str, Any]) -> List[tuple]:
        """从知识库检索相关内容"""
        # 这里应该连接到实际的知识库
        # 目前返回模拟数据
        
        max_results = task_data.get("max_results", self.default_max_results)
        
        # 模拟知识库检索结果
        mock_results = [
            ("机器学习是人工智能的一个分支，通过算法让计算机从数据中学习模式", 0.85),
            ("深度学习是机器学习的一种方法，使用神经网络模拟人脑学习过程", 0.80),
            ("监督学习需要标记数据来训练模型，无监督学习则不需要标记数据", 0.75),
            ("强化学习通过与环境交互来学习最优策略，常用于游戏和机器人控制", 0.70),
            ("Python是一种简单易学的编程语言，适合初学者", 0.65)
        ]
        
        # 根据查询内容返回相关结果
        if "python" in query.lower() or "编程" in query:
            return [mock_results[4]]  # Python相关
        elif "机器学习" in query or "machine learning" in query.lower():
            return mock_results[:3]  # 机器学习相关
        else:
            return mock_results[:max_results]  # 默认返回前几个
    
    def _calculate_confidence(self, retrieved_texts: List[tuple], threshold: float) -> float:
        """计算答案置信度"""
        if not retrieved_texts:
            return 0.0
        
        # 基于相似度分数计算平均置信度
        scores = [score for _, score in retrieved_texts]
        avg_score = sum(scores) / len(scores)
        
        # 如果平均分高于阈值，置信度较高
        if avg_score >= threshold:
            return min(avg_score, 0.95)
        else:
            return max(avg_score * 0.7, 0.1)
    
    async def handle_specialized_task(self, task_type: str, task_data: Dict[str, Any]) -> Dict[str, Any]:
        """处理特殊类型的问答任务"""
        
        if task_type == "fact_checking":
            return await self._fact_checking_task(task_data)
        elif task_type == "summary":
            return await self._summary_task(task_data)
        elif task_type == "comparison":
            return await self._comparison_task(task_data)
        else:
            # 默认处理
            return await self.process_task(task_data)
    
    async def _fact_checking_task(self, task_data: Dict[str, Any]) -> Dict[str, Any]:
        """事实核查任务"""
        statement = task_data.get("statement", "")
        
        # 修改查询为事实核查模式
        fact_check_query = f"请验证以下陈述是否正确：{statement}"
        
        modified_task_data = task_data.copy()
        modified_task_data["query"] = fact_check_query
        
        result = await self.process_task(modified_task_data)
        
        # 添加事实核查特定信息
        result["task_type"] = "fact_checking"
        result["statement"] = statement
        
        return result
    
    async def _summary_task(self, task_data: Dict[str, Any]) -> Dict[str, Any]:
        """摘要任务"""
        text = task_data.get("text", "")
        max_length = task_data.get("max_length", 100)
        
        summary_query = f"请将以下文本总结为不超过{max_length}字的摘要：{text}"
        
        modified_task_data = task_data.copy()
        modified_task_data["query"] = summary_query
        
        result = await self.process_task(modified_task_data)
        
        # 添加摘要特定信息
        result["task_type"] = "summary"
        result["original_text"] = text
        result["max_length"] = max_length
        
        return result
    
    async def _comparison_task(self, task_data: Dict[str, Any]) -> Dict[str, Any]:
        """比较任务"""
        item1 = task_data.get("item1", "")
        item2 = task_data.get("item2", "")
        
        comparison_query = f"请比较{item1}和{item2}的异同点"
        
        modified_task_data = task_data.copy()
        modified_task_data["query"] = comparison_query
        
        result = await self.process_task(modified_task_data)
        
        # 添加比较特定信息
        result["task_type"] = "comparison"
        result["item1"] = item1
        result["item2"] = item2
        
        return result